import math
import numpy as np


def TDMA_TridiagonalMatrix(a: np.ndarray, b: np.ndarray, c: np.ndarray, vM: np.ndarray):
    """
    使用高斯消元法解 三对角矩阵
    :param a: 三对角矩阵 a参数
    :param b: 三对角矩阵 b参数
    :param c: 三对角矩阵 c参数
    :param vM:  值矩阵
    :return np.ndarray :    解矩阵
    """
    # print(paraM.shape)
    length = vM.shape[0]  # 长度
    # 计算
    x = vM.copy()   # 返回结果
    c[0] = c[0] / b[0]
    x[0] = x[0] / b[0]
    for i in range(1, length):
        m = 1.0 / (b[i] - a[i] * c[i - 1])
        c[i] = c[i] * m
        x[i] = (x[i] - a[i] * x[i - 1]) * m

    for i in range(length - 2, -1 , -1):
        x[i] = x[i] - c[i] * x[i + 1]
    return x


def cubic_getval(anchor: np.ndarray, boundaryType=1):
    """
    计算3次样条取消各区间参数
    :param anchor: 锚点,2xn的矩阵 T[X, Y]
    :param boundaryType: 边界端点类型
    :return:
    """
    length = anchor.shape[1]    # 锚点数量
    # print(length)
    num = length - 1    # 区间数量
    x = anchor[0, :]
    y = anchor[1, :]

    # 计算步长
    h = np.zeros(length-1)
    for i in range(0, length-1):
        h[i] = x[i+1] - x[i]
    # 生成三对角 参数矩阵 与 值矩阵
    '''
    b c 0 ..... 0
    a b c ..    0
    0 a b c ... 0
    0 ..........0
    0 ........a b
    '''
    a = np.zeros(length)
    b = np.zeros(length)
    c = np.zeros(length)
    v = np.zeros(length)
    for i in range(1, length-1):
        # 参数
        a[i] = h[i-1]
        b[i] = 2 * (h[i-1] + h[i])
        c[i] = h[i]
        # 值
        v[i] = (y[i+1] - y[i])/h[i] - (y[i] - y[i-1])/h[i-1]
        v[i] = v[i] * 6

    # 边界值
    if boundaryType == 2:
        # 固定边界
        b[0] = 2 * h[0]
        c[0] = h[0]
        b[length-1] = 2 * h[length-2]
        a[length-1] = h[length-2]
    else:
        # 自由边界
        b[0] = 1
        v[0] = 0
        v[0] = 0

        b[length-1] = 1

    # 解方程
    m = TDMA_TridiagonalMatrix(a, b, c, v)

    # 计算参数
    arg = np.zeros((4, num))
    '''
    T[a,b,c,d]
    每个子区间有
    Gi(xRate) = ai + bi(xRate-xi) + ci(xRate-xi)**2 + di(xRate-xi)**3 
    '''
    for i in range(0, num):
        arg[0, i] = y[i]
        arg[1, i] = (y[i+1] - y[i])/h[i] - h[i]*m[i]/2 - h[i]*(m[i+1] - m[i])/6
        arg[2, i] = m[i]/2
        arg[3, i] = (m[i+1] - m[i])/h[i]/6

    return arg


def cal(arg, endpoint, x):
    return arg[0] +arg[1] * (x - endpoint) + arg[2] * (x - endpoint) ** 2 + arg[3] * (x - endpoint) ** 3


def generatorCubicSpline(anchor, spin):
    """
    根据锚点计算残次样条曲线数据
    返回曲线的数据[X, Y]
    X Y 为 1xn 的向量，X为取消X点，Y为曲线Y点，(xRate,yRate)组成曲线上的点，其中x每两个点之间的距离为spin
    :param anchor: 锚点 列表
    :param spin: X轴步进值
    :return: curveX, curveY
    """
    # 计算子区间 参数
    arg = cubic_getval(anchor, 1)
    num = arg.shape[1]  # 区间数量
    x = anchor[0, :]    # 锚点x向量
    pointNum = math.ceil((x[num] - x[0]) / spin) + 1  # 曲线总点数
    curveY = np.zeros(pointNum)  # 曲线y数据
    curveX = np.arange(pointNum) * spin  # 曲线x数据
    n = 0
    for i in range(0, num):
        cnt = int((x[i + 1] - x[i]) / spin)
        for j in range(0, cnt):
            curveY[n] = arg[0, i] + arg[1, i] * (j * spin) + arg[2, i] * (j * spin) ** 2 + arg[3, i] * (j * spin) ** 3
            # curve[n] = cal(arg[:, i], xRate[i], xRate[i] + j * rate)
            n = n + 1
    # 右端点
    curveY[pointNum - 1] = anchor[1, num]
    return curveX, curveY

# 测试

'''
# test TDMA
paraM = np.array(
    [
        [1,1,0],
        [1,2,1],
        [0,1,2]
    ]
)
vM = np.array(
    [3,4,5]
)
yRate = TDMA(paraM, vM)
print(yRate[:])
# should be [6, -3, 4]

'''


'''
import matplotlib.pyplot as plt

anchor = np.zeros((2, 4))
xRate = np.array([0, 0.3, 0.6, 1])
yRate = np.array([0, 0.4, 0.6, 1])
anchor[0, :] = xRate
anchor[1, :] = yRate

X, Y = generatorCubicSpline(anchor, 0.01)

plt.plot(X, Y,label='curve')

plt.show()
'''
'''

# plt.plot(xRate, yRate,label='sin')

arg = cubic_getval(anchor,1)
print(arg)
num = arg.shape[1]  # 区间数量
rate = 0.01  # 分辨率 相邻两点的x轴间隔为0.01
pointNum = math.ceil((xRate[num] - xRate[0])/0.01)  # 曲线总点数
curve = np.zeros(pointNum)
n = 0
for i in range(0, num):
    cnt = int((xRate[i+1] - xRate[i]) / rate)
    print('cnt',cnt)
    for j in range(0, cnt):
        curve[n] = cal(arg[:, i], xRate[i], xRate[i]+j*rate)
        n = n + 1
print(n)
xRate = np.arange(n) * rate
print('curve shape',curve.shape)
print('xRate shape', xRate.shape)
plt.plot(xRate, curve,label='curve')
plt.show()


'''